
# 程序旨在计算线性二次调节器控制器LQR, 计算反馈增益矩阵K

import math
import numpy as np
from scipy.linalg import solve_discrete_are
from scipy.linalg import solve_continuous_are



I1 = 3.31892 * 0.2225 * 0.2225
I2 = 0.09759 * 0.04 * 0.04
a = 3.31892 * 0.11125 * 0.11125 + 4 * 0.09759 * 0.11125 * 0.11125 + I1
b = (3.31892 + 2 * 0.09759) * 9.81 * 0.11125
c1 = 9.4e-3
c2 = 3.0e-4



A2_1 = b / a
A2_2 = -c1 / a
A2_4 = c2 / a
A4_1 = -b / a
A4_2 = c1 / a
A4_4 = -(a + I2) * c2 / (a * I2)

print(f"a={a}, b={b}, I1={I1}, I2={I2}")
print(f"A2_1={A2_1}, A2_2={A2_2}, A2_4={A2_4}, A4_1={A4_1}, A4_2={A4_2}, A4_4={A4_4}")

B2 = -1 / a
B4 = (a + I2) / (a * I2)

print(f"B2={B2}, B4={B4}")

def lqr_continuous(A, B, Q, R):
    """
    连续系统LQR控制器设计

    参数:
    A: 系统状态矩阵 (n x n)
    B: 输入矩阵 (n x m)
    Q: 状态权重矩阵 (n x n), 半正定
    R: 输入权重矩阵 (m x m), 正定

    返回:
    K: 最优反馈增益矩阵 (m x n)
    P: Riccati方程的解 (n x n)
    """

    # 求解连续时间代数Riccati方程
    P = solve_continuous_are(A, B, Q, R)

    # 计算最优反馈增益矩阵 K = R^{-1}B^T P
    K = np.linalg.inv(R) @ B.T @ P

    return K, P

def lqr_discrete(A, B, Q, R):
    """
    离散系统LQR控制器设计

    参数:
    A: 系统状态矩阵 (n x n)
    B: 输入矩阵 (n x m)
    Q: 状态权重矩阵 (n x n), 半正定
    R: 输入权重矩阵 (m x m), 正定

    返回:
    K: 最优反馈增益矩阵 (m x n)
    P: Riccati方程的解 (n x n)
    """

    # 求解离散时间代数Riccati方程
    P = solve_discrete_are(A, B, Q, R)

    # 计算最优反馈增益矩阵 K = (B^T P B + R)^{-1} B^T P A
    K = np.linalg.inv(B.T @ P @ B + R) @ B.T @ P @ A

    return K, P

# 示例使用
if __name__ == "__main__":
    # 定义系统矩阵（离散系统）
    A = np.array([[0, 1, 0, 0],
                  [A2_1, A2_2, 0, A2_4],
                  [0, 0, 0, 1],
                  [A4_1, A4_2, 0, A4_4]])
    B = np.array([[0],
                  [B2],
                  [0],
                  [B4]])

    # 定义权重矩阵
    Q = np.array([[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]])
    R = np.array([[1]])

    # 计算LQR增益
    K, P = lqr_discrete(A, B, Q, R)     # 离散系统
    #K, P = lqr_continuous(A, B, Q, R)  # 连续系统

    print("反馈增益矩阵 K:")
    print(K)
    print("\nRiccati方程解 P:")
    print(P)



